#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>

#include "../src/index/bw_tree.h"
#include "../src/memory/memory_manager.h"

// Test configuration
#define TEST_KEY_COUNT 1000
#define TEST_CONCURRENT_THREADS 4
#define TEST_STRESS_OPERATIONS 10000

// Test utilities
typedef struct {
    char key[64];
    char value[128];
} test_data_t;

// Generate test data
void generate_test_data(test_data_t* data, int count) {
    for (int i = 0; i < count; i++) {
        snprintf(data[i].key, sizeof(data[i].key), "key_%04d", i);
        snprintf(data[i].value, sizeof(data[i].value), "value_%04d_data", i);
    }
}

// Generate random test data
void generate_random_test_data(test_data_t* data, int count) {
    srand(time(NULL));
    for (int i = 0; i < count; i++) {
        snprintf(data[i].key, sizeof(data[i].key), "key_%04d", rand() % (count * 2));
        snprintf(data[i].value, sizeof(data[i].value), "value_%04d_random", i);
    }
}

// Test 1: Basic tree creation and destruction
void test_tree_creation() {
    printf("Running test_tree_creation...\n");
    
    memory_manager_t* mm = mm_init(1024 * 1024); // 1MB
    assert(mm != NULL);
    
    bw_tree_t* tree = bw_tree_create(mm);
    assert(tree != NULL);
    assert(tree->root_page_id != 0);
    assert(tree->mapping_table != NULL);
    assert(tree->gc_context != NULL);
    assert(tree->memory_mgr == mm);
    assert(atomic_load(&tree->total_nodes) == 1); // Root node
    assert(atomic_load(&tree->is_destroyed) == false);
    
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Tree creation and destruction test passed\n");
}

// Test 2: Basic insert and search operations
void test_basic_insert_search() {
    printf("Running test_basic_insert_search...\n");
    
    memory_manager_t* mm = mm_init(1024 * 1024);
    bw_tree_t* tree = bw_tree_create(mm);
    
    test_data_t data[10];
    generate_test_data(data, 10);
    
    // Insert test data
    for (int i = 0; i < 10; i++) {
        bw_result_t result = bw_tree_insert(tree, data[i].key, strlen(data[i].key),
                                           data[i].value, strlen(data[i].value));
        assert(result == BW_SUCCESS);
    }
    
    // Search for inserted data
    for (int i = 0; i < 10; i++) {
        bw_search_result_t search_result;
        bw_result_t result = bw_tree_search(tree, data[i].key, strlen(data[i].key), &search_result);
        
        assert(result == BW_SUCCESS);
        assert(search_result.found == true);
        assert(search_result.value != NULL);
        assert(search_result.value->size == strlen(data[i].value));
        assert(memcmp(search_result.value->data, data[i].value, strlen(data[i].value)) == 0);
        
        // Clean up search result
        if (search_result.value) {
            bw_value_destroy(search_result.value, mm);
        }
    }
    
    // Search for non-existent key
    bw_search_result_t search_result;
    bw_result_t result = bw_tree_search(tree, "non_existent", strlen("non_existent"), &search_result);
    assert(result == BW_ERROR_KEY_NOT_FOUND);
    assert(search_result.found == false);
    
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Basic insert and search test passed\n");
}

// Test 3: Delete operations
void test_delete_operations() {
    printf("Running test_delete_operations...\n");
    
    memory_manager_t* mm = mm_init(1024 * 1024);
    bw_tree_t* tree = bw_tree_create(mm);
    
    test_data_t data[10];
    generate_test_data(data, 10);
    
    // Insert test data
    for (int i = 0; i < 10; i++) {
        bw_tree_insert(tree, data[i].key, strlen(data[i].key),
                      data[i].value, strlen(data[i].value));
    }
    
    // Delete half of the data
    for (int i = 0; i < 5; i++) {
        bw_result_t result = bw_tree_delete(tree, data[i].key, strlen(data[i].key));
        assert(result == BW_SUCCESS);
    }
    
    // Verify deleted keys are not found
    for (int i = 0; i < 5; i++) {
        bw_search_result_t search_result;
        bw_result_t result = bw_tree_search(tree, data[i].key, strlen(data[i].key), &search_result);
        assert(result == BW_ERROR_KEY_NOT_FOUND);
    }
    
    // Verify remaining keys are still found
    for (int i = 5; i < 10; i++) {
        bw_search_result_t search_result;
        bw_result_t result = bw_tree_search(tree, data[i].key, strlen(data[i].key), &search_result);
        assert(result == BW_SUCCESS);
        assert(search_result.found == true);
        
        if (search_result.value) {
            bw_value_destroy(search_result.value, mm);
        }
    }
    
    // Test delete non-existent key
    bw_result_t result = bw_tree_delete(tree, "non_existent", strlen("non_existent"));
    assert(result == BW_ERROR_KEY_NOT_FOUND);
    
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Delete operations test passed\n");
}

// Test 4: Update operations
void test_update_operations() {
    printf("Running test_update_operations...\n");
    
    memory_manager_t* mm = mm_init(1024 * 1024);
    bw_tree_t* tree = bw_tree_create(mm);
    
    test_data_t data[5];
    generate_test_data(data, 5);
    
    // Insert test data
    for (int i = 0; i < 5; i++) {
        bw_tree_insert(tree, data[i].key, strlen(data[i].key),
                      data[i].value, strlen(data[i].value));
    }
    
    // Update values
    for (int i = 0; i < 5; i++) {
        char new_value[128];
        snprintf(new_value, sizeof(new_value), "updated_value_%d", i);
        
        bw_result_t result = bw_tree_update(tree, data[i].key, strlen(data[i].key),
                                           new_value, strlen(new_value));
        assert(result == BW_SUCCESS);
    }
    
    // Verify updated values
    for (int i = 0; i < 5; i++) {
        char expected_value[128];
        snprintf(expected_value, sizeof(expected_value), "updated_value_%d", i);
        
        bw_search_result_t search_result;
        bw_result_t result = bw_tree_search(tree, data[i].key, strlen(data[i].key), &search_result);
        
        assert(result == BW_SUCCESS);
        assert(search_result.found == true);
        assert(search_result.value != NULL);
        assert(search_result.value->size == strlen(expected_value));
        assert(memcmp(search_result.value->data, expected_value, strlen(expected_value)) == 0);
        
        if (search_result.value) {
            bw_value_destroy(search_result.value, mm);
        }
    }
    
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Update operations test passed\n");
}

// Test 5: Conditional operations
void test_conditional_operations() {
    printf("Running test_conditional_operations...\n");
    
    memory_manager_t* mm = mm_init(1024 * 1024);
    bw_tree_t* tree = bw_tree_create(mm);
    
    const char* key = "test_key";
    const char* value1 = "value1";
    const char* value2 = "value2";
    
    // Test insert_if_not_exists - should succeed
    bw_result_t result = bw_tree_insert_if_not_exists(tree, key, strlen(key), value1, strlen(value1));
    assert(result == BW_SUCCESS);
    
    // Test insert_if_not_exists - should fail (key exists)
    result = bw_tree_insert_if_not_exists(tree, key, strlen(key), value2, strlen(value2));
    assert(result == BW_ERROR_KEY_EXISTS);
    
    // Test update_if_exists - should succeed
    result = bw_tree_update_if_exists(tree, key, strlen(key), value2, strlen(value2));
    assert(result == BW_SUCCESS);
    
    // Verify the update worked
    bw_search_result_t search_result;
    result = bw_tree_search(tree, key, strlen(key), &search_result);
    assert(result == BW_SUCCESS);
    assert(memcmp(search_result.value->data, value2, strlen(value2)) == 0);
    if (search_result.value) {
        bw_value_destroy(search_result.value, mm);
    }
    
    // Test delete_if_exists - should succeed
    result = bw_tree_delete_if_exists(tree, key, strlen(key));
    assert(result == BW_SUCCESS);
    
    // Test update_if_exists on deleted key - should fail
    result = bw_tree_update_if_exists(tree, key, strlen(key), value1, strlen(value1));
    assert(result == BW_ERROR_KEY_NOT_FOUND);
    
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Conditional operations test passed\n");
}

// Test 6: Stress test with many operations
void test_stress_operations() {
    printf("Running test_stress_operations...\n");
    
    memory_manager_t* mm = mm_init(10 * 1024 * 1024); // 10MB
    bw_tree_t* tree = bw_tree_create(mm);
    
    test_data_t* data = malloc(TEST_KEY_COUNT * sizeof(test_data_t));
    generate_test_data(data, TEST_KEY_COUNT);
    
    // Insert many keys
    for (int i = 0; i < TEST_KEY_COUNT; i++) {
        bw_result_t result = bw_tree_insert(tree, data[i].key, strlen(data[i].key),
                                           data[i].value, strlen(data[i].value));
        assert(result == BW_SUCCESS);
        
        if (i % 100 == 0) {
            printf("  Inserted %d keys\n", i + 1);
        }
    }
    
    printf("  Inserted %d keys total\n", TEST_KEY_COUNT);
    
    // Verify all keys can be found
    int found_count = 0;
    for (int i = 0; i < TEST_KEY_COUNT; i++) {
        bw_search_result_t search_result;
        bw_result_t result = bw_tree_search(tree, data[i].key, strlen(data[i].key), &search_result);
        
        if (result == BW_SUCCESS && search_result.found) {
            found_count++;
            if (search_result.value) {
                bw_value_destroy(search_result.value, mm);
            }
        }
    }
    
    printf("  Found %d out of %d keys\n", found_count, TEST_KEY_COUNT);
    assert(found_count == TEST_KEY_COUNT);
    
    // Print statistics
    uint64_t total_nodes = atomic_load(&tree->total_nodes);
    uint64_t total_operations = atomic_load(&tree->total_operations);
    uint64_t consolidations = atomic_load(&tree->consolidations);
    
    printf("  Statistics: nodes=%llu, operations=%llu, consolidations=%llu\n", 
           (unsigned long long)total_nodes, (unsigned long long)total_operations, (unsigned long long)consolidations);
    
    free(data);
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Stress operations test passed\n");
}

// Thread data for concurrent testing
typedef struct {
    bw_tree_t* tree;
    memory_manager_t* mm;
    int thread_id;
    int start_key;
    int key_count;
    int* success_count;
    pthread_mutex_t* mutex;
} thread_data_t;

// Worker function for concurrent testing
void* concurrent_worker(void* arg) {
    thread_data_t* data = (thread_data_t*)arg;
    int local_success = 0;
    
    for (int i = 0; i < data->key_count; i++) {
        char key[64], value[128];
        snprintf(key, sizeof(key), "thread_%d_key_%d", data->thread_id, data->start_key + i);
        snprintf(value, sizeof(value), "thread_%d_value_%d", data->thread_id, data->start_key + i);
        
        bw_result_t result = bw_tree_insert(data->tree, key, strlen(key), value, strlen(value));
        if (result == BW_SUCCESS) {
            local_success++;
        }
        
        // Occasionally search for the key we just inserted
        if (i % 10 == 0) {
            bw_search_result_t search_result;
            bw_result_t search_res = bw_tree_search(data->tree, key, strlen(key), &search_result);
            if (search_res == BW_SUCCESS && search_result.value) {
                bw_value_destroy(search_result.value, data->mm);
            }
        }
    }
    
    pthread_mutex_lock(data->mutex);
    *data->success_count += local_success;
    pthread_mutex_unlock(data->mutex);
    
    return NULL;
}

// Test 7: Concurrent operations
void test_concurrent_operations() {
    printf("Running test_concurrent_operations...\n");
    
    memory_manager_t* mm = mm_init(20 * 1024 * 1024); // 20MB
    bw_tree_t* tree = bw_tree_create(mm);
    
    pthread_t threads[TEST_CONCURRENT_THREADS];
    thread_data_t thread_data[TEST_CONCURRENT_THREADS];
    int success_count = 0;
    pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
    
    int keys_per_thread = TEST_KEY_COUNT / TEST_CONCURRENT_THREADS;
    
    // Create threads
    for (int i = 0; i < TEST_CONCURRENT_THREADS; i++) {
        thread_data[i].tree = tree;
        thread_data[i].mm = mm;
        thread_data[i].thread_id = i;
        thread_data[i].start_key = i * keys_per_thread;
        thread_data[i].key_count = keys_per_thread;
        thread_data[i].success_count = &success_count;
        thread_data[i].mutex = &mutex;
        
        int result = pthread_create(&threads[i], NULL, concurrent_worker, &thread_data[i]);
        assert(result == 0);
    }
    
    // Wait for all threads to complete
    for (int i = 0; i < TEST_CONCURRENT_THREADS; i++) {
        pthread_join(threads[i], NULL);
    }
    
    printf("  Concurrent operations: %d successful inserts out of %d attempted\n", 
           success_count, TEST_KEY_COUNT);
    
    // Print final statistics
    uint64_t total_nodes = atomic_load(&tree->total_nodes);
    uint64_t total_operations = atomic_load(&tree->total_operations);
    uint64_t consolidations = atomic_load(&tree->consolidations);
    
    printf("  Final statistics: nodes=%llu, operations=%llu, consolidations=%llu\n", 
           (unsigned long long)total_nodes, (unsigned long long)total_operations, (unsigned long long)consolidations);
    
    pthread_mutex_destroy(&mutex);
    bw_tree_destroy(tree);
    mm_cleanup(mm);
    
    printf("✓ Concurrent operations test passed\n");
}

// Test 8: Memory management and cleanup
void test_memory_management() {
    printf("Running test_memory_management...\n");
    
    memory_manager_t* mm = mm_init(1024 * 1024);
    size_t initial_allocated = mm->total_allocated;
    
    {
        bw_tree_t* tree = bw_tree_create(mm);
        
        // Insert some data
        for (int i = 0; i < 100; i++) {
            char key[32], value[64];
            snprintf(key, sizeof(key), "key_%d", i);
            snprintf(value, sizeof(value), "value_%d", i);
            bw_tree_insert(tree, key, strlen(key), value, strlen(value));
        }
        
        size_t peak_allocated = mm->total_allocated;
        printf("  Peak memory usage: %zu bytes\n", peak_allocated);
        
        bw_tree_destroy(tree);
    }
    
    // Give some time for garbage collection
    sleep(1);
    
    size_t final_allocated = mm->total_allocated;
    printf("  Memory after cleanup: %zu bytes (initial: %zu)\n", final_allocated, initial_allocated);
    
    mm_cleanup(mm);
    
    printf("✓ Memory management test passed\n");
}

// Main test runner
int main() {
    printf("Starting Bw-Tree Test Suite\n");
    printf("===========================\n\n");
    
    // Run all tests
    test_tree_creation();
    test_basic_insert_search();
    test_delete_operations();
    test_update_operations();
    test_conditional_operations();
    test_stress_operations();
    test_concurrent_operations();
    test_memory_management();
    
    printf("\n===========================\n");
    printf("All Bw-Tree tests passed! ✓\n");
    
    return 0;
}